package org.crazycake.shiro;

import redis.clients.jedis.*;
import redis.clients.jedis.exceptions.JedisConnectionException;
import redis.clients.jedis.params.ScanParams;

import java.util.*;

public class RedisClusterManager implements IRedisManager {

    private static final int DEFAULT_COUNT = 100;
    private static final int DEFAULT_MAX_ATTEMPTS = 3;
    private static final String DEFAULT_HOST = "127.0.0.1:7000,127.0.0.1:7001,127.0.0.1:7002";
    private String host = DEFAULT_HOST;

    // timeout for jedis try to connect to redis server, not expire time! In milliseconds
    private int timeout = Protocol.DEFAULT_TIMEOUT;

    // timeout for jedis try to read data from redis server
    private int soTimeout = Protocol.DEFAULT_TIMEOUT;

    private String password;

    private int database = Protocol.DEFAULT_DATABASE;

    private int count = DEFAULT_COUNT;

    private int maxAttempts = DEFAULT_MAX_ATTEMPTS;

    /**
     * JedisPoolConfig used to initialize JedisPool.
     */
    private JedisPoolConfig jedisPoolConfig = new JedisPoolConfig();

    private volatile JedisCluster jedisCluster = null;

    private void init() {
        if (jedisCluster == null) {
            synchronized (RedisClusterManager.class) {
                if (jedisCluster == null) {
                    JedisClientConfig clientConfig = DefaultJedisClientConfig.builder()
                            .password(password) // 密码
                            .socketTimeoutMillis(timeout) // Socket超时（读取操作）
                            .connectionTimeoutMillis(soTimeout) // 连接超时
                            .clientName("myApp") // 客户端名称
                            .build();

                    jedisCluster = new JedisCluster(getHostAndPortSet(), clientConfig);
                }
            }
        }
    }

    private Set<HostAndPort> getHostAndPortSet() {
        String[] hostAndPortArr = host.split(",");
        Set<HostAndPort> hostAndPorts = new HashSet<HostAndPort>();
        for (String hostAndPortStr : hostAndPortArr) {
            String[] hostAndPort = hostAndPortStr.split(":");
            hostAndPorts.add(new HostAndPort(hostAndPort[0], Integer.parseInt(hostAndPort[1])));
        }
        return hostAndPorts;
    }


    protected JedisCluster getJedisCluster() {
        if (jedisCluster == null) {
            init();
        }
        return jedisCluster;
    }

    @Override
    public byte[] get(byte[] key) {
        if (key == null) {
            return null;
        }
        return getJedisCluster().get(key);
    }

    @Override
    public byte[] set(byte[] key, byte[] value, int expireTime) {
        if (key == null) {
            return null;
        }
        getJedisCluster().set(key, value);
        if (expireTime >= 0) {
            getJedisCluster().expire(key, expireTime);
        }
        return value;
    }

    @Override
    public void del(byte[] key) {
        if (key == null) {
            return;
        }
        getJedisCluster().del(key);
    }

    @Override
    public Long dbSize(byte[] pattern) {
        long totalSize = 0L;
        // 获取集群所有节点连接池（新版返回Map<String, ConnectionPool>）
        Map<String, ConnectionPool> clusterNodes = jedisCluster.getClusterNodes();

        for (Map.Entry<String, ConnectionPool> entry : clusterNodes.entrySet()) {
            // 跳过从节点，只统计主节点数据
            if (isReplica(entry.getValue())) continue;

            try (Connection connection = entry.getValue().getResource()) {
                // 新版Jedis使用ProtocolCommand替代字符串命令
                connection.sendCommand(Protocol.Command.DBSIZE);
                Long nodeSize = connection.getIntegerReply();
                if (nodeSize != null) {
                    totalSize += nodeSize;
                }
            } catch (JedisConnectionException e) {
                // 节点不可用时不中断统计，但记录日志
                System.err.printf("Node %s unreachable: %s%n", entry.getKey(), e.getMessage());
            }
        }

        return totalSize;
    }

    /**
     * 判断节点是否为从节点
     */
    private boolean isReplica(ConnectionPool pool) {
        try (Connection connection = pool.getResource()) {
            connection.sendCommand(Protocol.Command.INFO, "replication".getBytes());
            String info = connection.getBulkReply();
            return info.contains("role:slave");
        } catch (Exception e) {
            return false; // 默认为主节点
        }
    }




    /**
     * 获取集群中所有匹配模式的键（生产环境慎用！）
     *
     * @param pattern key的模式（如：*user*）
     * @return 匹配的所有键
     */
    /**
     * 获取集群中匹配模式的所有键
     */
    @Override
    public Set<byte[]> keys(byte[] pattern) {
        Set<byte[]> keys = new HashSet<>();
        Map<String, ConnectionPool> clusterNodes = jedisCluster.getClusterNodes();

        for (Map.Entry<String, ConnectionPool> node : clusterNodes.entrySet()) {
            if (isReplica(node.getValue())) continue;

            try (Connection conn = node.getValue().getResource()) {
                keys.addAll(scanAllKeys(conn, pattern));
            } catch (JedisConnectionException e) {
                System.err.println("节点连接失败: " + node.getKey());
            }
        }

        return keys;
    }

    /**
     * 安全扫描所有键（生产推荐）
     */
    public void scanKeys(byte[] pattern, int batchSize, Consumer<byte[]> consumer) {
        for (ConnectionPool pool : getMasterPools()) {
            try (Connection conn = pool.getResource()) {
                scanKeysInNode(conn, pattern, batchSize, consumer);
            }
        }
    }

    // 私有工具方法
    private Set<byte[]> scanAllKeys(Connection conn, byte[] pattern) {
        Set<byte[]> keys = new HashSet<>();
        String cursor = "0";

        do {
            cursor = scanBatch(conn, cursor, pattern, 100, keys);
        } while (!"0".equals(cursor));

        return keys;
    }

    private String scanBatch(Connection conn, String cursor, byte[] pattern, int count, Set<byte[]> keys) {
        List<byte[]> args = new ArrayList<>();
        args.add(cursor.getBytes());

        if (pattern != null) {
            args.add(Protocol.Keyword.MATCH.getRaw());
            args.add(pattern);
        }

        args.add(Protocol.Keyword.COUNT.getRaw());
        args.add(Protocol.toByteArray(count));

        conn.sendCommand(Protocol.Command.SCAN, args.toArray(new byte[0][]));
        List<Object> response = conn.getObjectMultiBulkReply();

        String newCursor = new String((byte[]) response.get(0));
        keys.addAll((List<byte[]>) response.get(1));

        return newCursor;
    }

    private void scanKeysInNode(Connection conn, byte[] pattern, int batchSize, Consumer<byte[]> consumer) {
        String cursor = "0";

        do {
            Set<byte[]> batch = new HashSet<>();
            cursor = scanBatch(conn, cursor, pattern, batchSize, batch);

            for (byte[] key : batch) {
                consumer.accept(key);
            }
        } while (!"0".equals(cursor));
    }

    private List<ConnectionPool> getMasterPools() {
        List<ConnectionPool> masters = new ArrayList<>();
        Map<String, ConnectionPool> allNodes = jedisCluster.getClusterNodes();

        for (ConnectionPool pool : allNodes.values()) {
            if (!isReplica(pool)) {
                masters.add(pool);
            }
        }

        return masters;
    }



    // 键消费函数接口
    @FunctionalInterface
    public interface Consumer<T> {
        void accept(T t);
    }



    private Set<byte[]> getKeysFromClusterNode(ConnectionPool pool, byte[] pattern) {
        Set<byte[]> keys = new HashSet<>();

        // 使用 try-with-resources 确保连接自动关闭
        try (Connection connection = pool.getResource()) {
            String cursor = ScanParams.SCAN_POINTER_START;

            do {
                // 准备 SCAN 命令参数
                List<byte[]> commandArgs = new ArrayList<>();
                commandArgs.add(cursor.getBytes());  // 当前游标

                if (pattern != null && pattern.length > 0) {
                    commandArgs.add(Protocol.Keyword.MATCH.getRaw());  // MATCH 参数
                    commandArgs.add(pattern);                // 模式匹配
                }

                commandArgs.add(Protocol.Keyword.COUNT.getRaw());      // COUNT 参数
                commandArgs.add(Protocol.toByteArray(100));  // 每批处理100个键

                // 发送 SCAN 命令
                connection.sendCommand(Protocol.Command.SCAN,
                        commandArgs.toArray(new byte[commandArgs.size()][]));

                // 解析响应：第一个元素是新游标，第二个元素是匹配的键列表
                List<Object> response = connection.getObjectMultiBulkReply();

                // 更新游标
                cursor = new String((byte[]) response.get(0));

                // 添加匹配的键到结果集
                List<byte[]> resultKeys = (List<byte[]>) response.get(1);
                keys.addAll(resultKeys);

            } while (!ScanParams.SCAN_POINTER_START.equals(cursor)); // 直到游标返回0

        } catch (JedisConnectionException e) {
            System.err.println("节点连接异常: " + e.getMessage());
            // 添加重试逻辑或日志记录
        }

        return keys;
    }







    private long getDbSizeFromClusterNode(ConnectionPool pool, byte[] pattern) {
        long totalKeys = 0L;

        // 使用 try-with-resources 确保连接正确关闭
        try (Connection connection = pool.getResource()) {
            String cursor = ScanParams.SCAN_POINTER_START;

            do {
                // 构建 SCAN 命令参数
                List<byte[]> commandArgs = new ArrayList<>();
                commandArgs.add(cursor.getBytes());  // 当前游标

                // 添加匹配模式（如果提供）
                if (pattern != null && pattern.length > 0) {
                    commandArgs.add(Protocol.Keyword.MATCH.getRaw());  // MATCH 参数
                    commandArgs.add(pattern);                 // 匹配模式
                }

                commandArgs.add(Protocol.Keyword.COUNT.getRaw());      // COUNT 参数
                commandArgs.add(Protocol.toByteArray(500));   // 每批处理500个键

                // 发送 SCAN 命令
                connection.sendCommand(Protocol.Command.SCAN,
                        commandArgs.toArray(new byte[commandArgs.size()][]));

                // 解析响应：第一个元素是新游标，第二个元素是匹配的键列表
                List<Object> response = connection.getObjectMultiBulkReply();

                // 更新游标
                cursor = new String((byte[]) response.get(0));

                // 获取匹配的键列表并统计数量（关键修复）
                List<byte[]> matchedKeys = (List<byte[]>) response.get(1);
                totalKeys += matchedKeys.size();  // 累加实际匹配的键数量

            } while (!ScanParams.SCAN_POINTER_START.equals(cursor)); // 直到游标返回0

        } catch (JedisConnectionException e) {
            System.err.println("节点连接失败: " + e.getMessage());
            // 生产环境应记录日志或告警
        }

        return totalKeys;
    }





    public int getMaxAttempts() {
        return maxAttempts;
    }

    public void setMaxAttempts(int maxAttempts) {
        this.maxAttempts = maxAttempts;
    }

    public String getHost() {
        return host;
    }

    public void setHost(String host) {
        this.host = host;
    }

    public int getTimeout() {
        return timeout;
    }

    public void setTimeout(int timeout) {
        this.timeout = timeout;
    }

    public int getSoTimeout() {
        return soTimeout;
    }

    public void setSoTimeout(int soTimeout) {
        this.soTimeout = soTimeout;
    }

    public String getPassword() {
        return password;
    }

    public void setPassword(String password) {
        this.password = password;
    }

    public int getDatabase() {
        return database;
    }

    public void setDatabase(int database) {
        this.database = database;
    }

    public int getCount() {
        return count;
    }

    public void setCount(int count) {
        this.count = count;
    }

    public void setJedisCluster(JedisCluster jedisCluster) {
        this.jedisCluster = jedisCluster;
    }

    public JedisPoolConfig getJedisPoolConfig() {
        return jedisPoolConfig;
    }

    public void setJedisPoolConfig(JedisPoolConfig jedisPoolConfig) {
        this.jedisPoolConfig = jedisPoolConfig;
    }
}
